@echo off
chcp 65001
setlocal enabledelayedexpansion

echo ========================================
echo           LoRA 训练脚本
echo ========================================
echo.

:: 设置Python路径
set PYTHON_PATH=E:\AI\python\python.exe

:: 设置脚本路径
set SCRIPT_PATH=%~dp0training.py
echo 正在调用脚本: %SCRIPT_PATH%

:: 执行模式选择
echo ========================================
echo           选择执行模式
echo ========================================
echo 1. thread     - 线程模式 (推荐)
echo    在当前进程中使用线程执行训练
echo    资源共享，速度较快，适合大多数情况
echo.
echo 2. subprocess - 子进程模式
echo    创建独立子进程执行训练  
echo    资源隔离，更稳定，适合长时间训练
echo.
echo ========================================

:choose_mode
set /p mode_choice="请输入选择 (1 或 2，直接回车默认选择1): "

if "%mode_choice%"=="" set mode_choice=1

if "%mode_choice%"=="1" (
set EXECUTION_MODE=thread
    echo.
    echo 已选择: 线程模式 ^(thread^)
    goto mode_selected
)

if "%mode_choice%"=="2" (
    set EXECUTION_MODE=subprocess
    echo.
    echo 已选择: 子进程模式 ^(subprocess^)
    goto mode_selected
)

echo.
echo 无效选择，请输入 1 或 2
echo.
goto choose_mode

:mode_selected
echo.
set MODEL_PATH=E:\AI\ComfyUI\models\checkpoints\v1-5-pruned-emaonly-fp16.safetensors
set DATA_DIR=E:\training_data\images
set OUTPUT_DIR=E:\training_data\LoraOutput

:: 验证训练数据目录
if not exist "%DATA_DIR%" (
    echo 错误: 训练数据目录不存在: %DATA_DIR%
    pause
    exit /b 1
)


set OUTPUT_NAME=my_lora
set NETWORK_DIM=32
set NETWORK_ALPHA=32
set BATCH_SIZE=1
set LR=0.0001
set STEPS=10
set RESOLUTION=512

:: 设置 PYTHONPATH
set PYTHONPATH=E:\AI\ComfyUI\custom_nodes\comfyui_lora_trainer;E:\AI\ComfyUI\custom_nodes\comfyui_lora_trainer\sd_scripts;%PYTHONPATH%
echo PYTHONPATH设置为: %PYTHONPATH%
echo.
echo ========================================
echo 开始训练，执行模式: %EXECUTION_MODE%
echo ========================================

:: 执行命令
%PYTHON_PATH% "%SCRIPT_PATH%" ^
    --execution_mode %EXECUTION_MODE% ^
    --pretrained_model_name_or_path "%MODEL_PATH%" ^
    --train_data_dir "%DATA_DIR%" ^
    --output_dir "%OUTPUT_DIR%" ^
    --output_name "%OUTPUT_NAME%" ^
    --network_dim %NETWORK_DIM% ^
    --network_alpha %NETWORK_ALPHA% ^
    --train_batch_size %BATCH_SIZE% ^
    --learning_rate %LR% ^
    --max_train_steps %STEPS% ^
    --save_model_as safetensors ^
    --mixed_precision fp16 ^
    --save_every_n_steps 500 ^
    --seed 42 ^
    --resolution %RESOLUTION%,%RESOLUTION% ^
    --caption_extension .txt ^
    --network_module networks.lora ^
    --bucket_reso_steps=64 ^
    --min_bucket_reso=320 ^
    --max_bucket_reso=1280 ^
    --enable_bucket ^
    --logging_dir "%OUTPUT_DIR%\logs" ^
    --cache_latents ^
    --xformers ^
    --gradient_checkpointing ^
    --cache_latents_to_disk ^
    --gradient_accumulation_steps 1



if %ERRORLEVEL% neq 0 (
    echo 训练失败，错误代码: %ERRORLEVEL%
    pause
    exit /b %ERRORLEVEL%
)

echo 训练完成!
pause
